package com.yww.leetcode.binarytree;

import com.yww.leetcode.binarytree.common.TreeNode;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.function.Consumer;

/**
 * @author yww
 * @description
 * @since 2024/2/8 16:51
 */
public class Traversal144 {
    public List<Integer> preorderTraversal(TreeNode root) {
        List<Integer> list = new ArrayList<>();
        traversal(root, list::add, null, null);
        return list;
    }

    public List<Integer> inorderTraversal(TreeNode root) {
        List<Integer> list = new ArrayList<>();
        traversal(root, null, list::add, null);
        return list;
    }

    public List<Integer> postorderTraversal(TreeNode root) {
        List<Integer> list = new ArrayList<>();
        traversal(root, null, null, list::add);
        return list;
    }


    /**
     * 迭代实现的二叉树遍历
     *
     * @param root      根节点
     * @param preOrder  前序遍历消费器
     * @param inOrder   中序遍历消费器
     * @param postOrder 后续遍历消费器
     **/
    public static void traversal(TreeNode root,
                                 Consumer<Integer> preOrder,
                                 Consumer<Integer> inOrder,
                                 Consumer<Integer> postOrder) {
        TreeNode cur = root;
        LinkedList<TreeNode> stack = new LinkedList<>(); // 记录待返回的节点栈
        TreeNode poped = null; // 记录弹出的节点
        while (!stack.isEmpty() || cur != null) {
            if (cur != null) { // 处理当前节点
                stack.push(cur); // 入栈
                if (preOrder != null) {
                    preOrder.accept(cur.val); // 处理左子树之前，进行前序消费
                }
                cur = cur.left; // 处理左子树
            } else { // 需要返回
                TreeNode peek = stack.peek();
                if (peek.right == null) { // 右子树空，无需处理
                    if (inOrder != null) {
                        inOrder.accept(peek.val); // 弹出之前，进行中序消费
                    }
                    poped = stack.pop(); // 弹出栈顶
                    if (postOrder != null) {
                        postOrder.accept(peek.val); // 弹出后，最后进行后序消费
                    }
                } else if (peek.right == poped) { // 右子树已经处理完毕，当前节点直接弹出
                    poped = stack.pop();
                    if (postOrder != null) {
                        postOrder.accept(peek.val); // 弹出后，最后进行后序消费
                    }
                } else { // 处理右子树
                    if (inOrder != null) {
                        inOrder.accept(peek.val); // 处理右子树之前，进行中序消费
                    }
                    cur = peek.right;
                }
            }
        }
    }
}
